import numpy as np
import os
import re
from DDPM.Evaluation.InternalBallistic import get_target_from_generated_samples,get_av1_values_from_generated_samples, get_av2_values_from_generated_samples
import DEAD.AutoDecoder.Evaluation.InternalBallistic as IB
from Utility.PlotCurve import plot_curves_optimization,plot_curves_optimization_no_legend
import torch
from OPT.Objective import performance_matching_degree

def start_internal_ballistic_cal(control_params):
    result_path=control_params["result_path"]
    t_array_target, p_array_target = get_target_from_generated_samples(result_path)
    cstar, rho, n, At = get_av2_values_from_generated_samples(result_path)
    cstar_rho = cstar * rho

     # Get all files in the directory
    files = os.listdir(f"{result_path}")
    # 过滤出 .npz 结尾的文件
    npz_files = [file for file in files if file.endswith('.npz') and re.search('#', file)]
    npz_files = sorted(npz_files, key=lambda x: int(re.search(r'(\d+)#', x).group(1)))

    print(npz_files)

    t_array_list = []
    p_array_list = []
    r2_list=[]
    r2_k_list=[]
    for i in [0,-1]:
        index_in_samples = int(re.search(r'(\d+)#', npz_files[i]).group(1))
        r_ref, R = get_av1_values_from_generated_samples(result_path, index_in_samples)
        npz_file = np.load(f"{result_path}/{npz_files[i]}")
        w_array_norm = npz_file['burned_web_array']
        Ab_array_norm = npz_file['burning_surface_area_array']
        t_array_designed, p_array_designed = IB.internal_ballistic_calculation(
            w_array_norm=w_array_norm,
            Ab_array_norm=Ab_array_norm,
            cstar_rho=cstar_rho,
            r_ref=r_ref,
            n=n,
            R=R,
            At=At,
            verify_range=False)
        t_array_list.append(t_array_designed)
        p_array_list.append(p_array_designed)

        t_array_target_tensor = torch.tensor(t_array_target, dtype=torch.float32)
        p_array_target_tensor = torch.tensor(p_array_target, dtype=torch.float32)
        t_array_designed_tensor = torch.tensor(t_array_designed, dtype=torch.float32)
        p_array_designed_tensor = torch.tensor(p_array_designed, dtype=torch.float32)

        with torch.no_grad():
            r2, r2_k, t_vector_designed, p_vector_designed=\
                performance_matching_degree(t_array_designed_tensor,
                                        p_array_designed_tensor,
                                        t_array_target_tensor,
                                        p_array_target_tensor)
            r2_list.append(r2.item())
            r2_k_list.append(r2_k.item())
    np.savez_compressed(f"{result_path}/internal_ballistic.npz",
                    t_array_design=t_array_list,
                    p_array_design=p_array_list,
                    t_arstart_internal_ballistic_calray_target=t_array_target,
                    p_array_target=p_array_target,
                    r2_k_list=r2_k_list,
                    r2_list=r2_list,
                    r2_mean=np.mean(r2_list))
    
    plot_curves_optimization(title='Chamber Pressure vs Burned time',
                xlabel='$\it{t}$ /s',  # t 设置为斜体，保留字体
                ylabel='$\it{p}$ /MPa',  # p 设置为斜体，保留字体
                x=np.array(t_array_list).T,
                y=np.array(p_array_list).T*1e-6,
                ref_x=t_array_target,
                ref_y=p_array_target*1e-6,
                filename=f"{result_path}/Optimization.svg")

    plot_curves_optimization_no_legend(title='Chamber Pressure vs Burned time',
                xlabel='$\it{t}$ /s',  # t 设置为斜体，保留字体
                ylabel='$\it{p}$ /MPa',  # p 设置为斜体，保留字体
                x=np.array(t_array_list).T,
                y=np.array(p_array_list).T*1e-6,
                ref_x=t_array_target,
                ref_y=p_array_target*1e-6,
                filename=f"{result_path}/Optimization_no_legend.svg")

if __name__ == "__main__":
    from OPT.Config import *
    define =make_define("dicA",1,1e-3,1)
    control_params, dead_dic = get_control_params(define)
    start_internal_ballistic_cal(control_params)